import os
from PIL import Image

# --- 配置参数 ---

# 1. 输入目录：包含子目录和图片的根文件夹
INPUT_DIR = 'images'

# 2. 输出目录：用于保存处理后图片的根文件夹（如果不存在，脚本会自动创建）
OUTPUT_DIR = 'image_resized'

# 3. 目标尺寸：你希望所有图片最终变成的尺寸 (宽度, 高度)
TARGET_SIZE = (768, 768)

# 4. 填充颜色：使用RGB元组。(0, 0, 0)是黑色, (255, 255, 255)是白色
PADDING_COLOR = (255, 255, 255) 

# --- 核心处理函数 ---

def resize_with_padding(img, target_size, padding_color):
    """
    按比例缩放图片并填充到目标尺寸。

    :param img: PIL.Image 对象
    :param target_size: 目标尺寸 (宽度, 高度)
    :param padding_color: 填充颜色 (R, G, B)
    :return: 处理后的 PIL.Image 对象
    """
    # 原始图片的尺寸
    original_width, original_height = img.size

    # 目标尺寸
    target_width, target_height = target_size

    # 计算缩放比例，以保持宽高比
    # 取宽度缩放比和高度缩放比中较小的一个，确保整个图片能放进目标框
    ratio = min(target_width / original_width, target_height / original_height)

    # 根据比例计算缩放后的新尺寸
    new_width = int(original_width * ratio)
    new_height = int(original_height * ratio)

    # 使用高质量的LANCZOS算法进行缩放
    resized_img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    # 创建一个新的背景图（画布），使用指定的填充颜色
    # 确保图片是 'RGB' 模式，以支持彩色填充
    new_img = Image.new('RGB', target_size, padding_color)
    
    # 计算粘贴位置，使其居中
    paste_x = (target_width - new_width) // 2
    paste_y = (target_height - new_height) // 2
    
    # 将缩放后的图片粘贴到背景图的中心
    new_img.paste(resized_img, (paste_x, paste_y))
    
    return new_img

# --- 主程序 ---

def main():
    """
    遍历目录，处理所有图片。
    """
    # 检查输入目录是否存在
    if not os.path.isdir(INPUT_DIR):
        print(f"错误：输入目录 '{INPUT_DIR}' 不存在。")
        return

    # 遍历输入目录下的所有文件和子目录
    for root, _, files in os.walk(INPUT_DIR):
        # 创建对应的输出子目录结构
        # os.path.relpath 计算从INPUT_DIR到当前root的相对路径
        # 例如，将 'image/class_A' 变成 'class_A'
        relative_path = os.path.relpath(root, INPUT_DIR)
        
        # 如果是根目录本身，relative_path会是'.'，我们不希望创建名为'.'的文件夹
        if relative_path == '.':
            output_subdir = OUTPUT_DIR
        else:
            output_subdir = os.path.join(OUTPUT_DIR, relative_path)

        # 创建输出子目录（如果不存在）
        os.makedirs(output_subdir, exist_ok=True)
        
        # 过滤并处理图片文件
        for filename in files:
            # 通过文件扩展名判断是否为图片
            if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
                input_path = os.path.join(root, filename)
                output_path = os.path.join(output_subdir, filename)
                
                try:
                    # 打开图片
                    with Image.open(input_path) as img:
                        # 转换图片为RGB模式，以处理PNG等带透明通道或不同模式的图片
                        img_rgb = img.convert('RGB')
                        
                        # 调用函数进行缩放和填充
                        processed_img = resize_with_padding(img_rgb, TARGET_SIZE, PADDING_COLOR)
                        
                        # 保存处理后的图片
                        processed_img.save(output_path)
                        print(f"成功处理: {input_path} -> {output_path}")

                except Exception as e:
                    print(f"处理失败: {input_path} | 错误: {e}")
                    with open('error.log', 'a') as f:
                        f.write(f"{input_path}\n")

    print("\n所有图片处理完成！")


if __name__ == '__main__':
    main()